#!/usr/bin/env python3

from pathlib import Path
PATH_ROOT = Path(__file__).parents[4]
PATH_PARENT = Path(__file__).parent
PATH_LIGHTGLUE = PATH_ROOT / "pkgs/LightGlue-ONNX"
import sys
sys.path.append(str(PATH_LIGHTGLUE))
from map_aligner import MapAligner

import numpy as np
import cv2
import matplotlib.pyplot as plt

def process_img(path_img: Path):
    map = cv2.imread(str(path_img))
    height, width = map.shape[:2]
    x, y = path_img.stem.split('_')[-2:]
    x, y = int(float(x)/10), - height - int(float(y)/10)
    A = np.array([[1, 0, x], [0, 1, y]])
    resolution = 0.1
    return map, A, resolution

def draw_point(img: np.ndarray, point: tuple, color=(0, 255, 0), radius=2):
    return cv2.circle(img, (int(point[0]), int(point[1])), radius, color, -1)

def adjust_map(path_origin: Path, path_camera: Path):
    device = 'cuda'
    map_aligner = MapAligner(device=device)
    delta_angle = 36.0
    score_threshold = 80

    origin_map, origin_A, origin_resolution = process_img(path_origin)
    camera_map, camera_A, camera_resolution = process_img(path_camera)
    assert camera_resolution == origin_resolution, f"Camera map resolution = {camera_resolution:.2f} must match origin map resolution = {origin_resolution:.2f}."

    com_A, score = map_aligner.rotate_compare(
        camera_map, origin_map, delta_angle=delta_angle,
        verbose=True, save=False
    )
    if score < score_threshold:
        print(f'Not enough matching points: {score} < {score_threshold}, skipping pose adjustment.')
        return
    if com_A is None:
        return
    stack = lambda A: np.vstack([A, [0, 0, 1]])
    A = origin_A @ stack(com_A) @ np.linalg.inv(stack(camera_A))
    print(f"A={A}")


    corigin = np.r_[np.zeros((2, 1)), np.ones((1, 1))]
    oorigin2 = (np.linalg.inv(stack(origin_A)) @ stack(A) @ corigin)[:2, 0]
    oorigin3 = A @ corigin
    print(oorigin2, oorigin3)

    cimg = camera_map.copy()
    # cimg = draw_point(cimg, -camera_A[:2, 2])
    oimg = origin_map.copy()
    # oimg = draw_point(oimg, -origin_A[:2, 2])
    # oimg = draw_point(oimg, oorigin2, color=(255, 0, 0))

    fig, axs = plt.subplots(1, 2, figsize=(14, 8))
    axs[0].imshow(cimg)
    axs[0].set_title('Camera Map')
    axs[0].scatter(-camera_A[0, 2], -camera_A[1, 2], color='blue', s=50, label='Camera Origin')
    axs[0].legend()
    axs[1].imshow(oimg)
    axs[1].set_title('Origin Map')
    axs[1].scatter(-origin_A[0, 2], -origin_A[1, 2], color='blue', s=50, label='Origin')
    axs[1].scatter(oorigin2[0], oorigin2[1], color='green', s=50, label='Adjusted Camera Origin')
    axs[1].legend()
    fig.tight_layout()
    plt.show()

if __name__ == '__main__':
    path_origin = PATH_PARENT / "assets/2d_map_-1867_-811.png"
    path_camera = PATH_PARENT / "assets/2d_map_-730_-1722.png"
    adjust_map(path_origin, path_camera)
